import random

import torch
import pandas as pd
import csv
import os
import cv2
from Config import Config
import matplotlib.pyplot as plt
import timm
from simdataset import Simdataset
from sklearn.model_selection import KFold
import albumentations as A
from albumentations.pytorch import ToTensorV2
from torch.utils.data import DataLoader,Dataset
from sklearn.metrics import roc_auc_score
import numpy as np
import glob
img1_transformer = A.Compose([
    A.Resize(height=Config.img_size, width=Config.img_size),
    A.Normalize(max_pixel_value=255.0, p=1.0),
    ToTensorV2(p=1.0),
])
class Simdataset(Dataset):
    def __init__(self,path,img1_transformer=img1_transformer,img2_transformer=img1_transformer):
        self.path = path
        self.data = pd.read_csv(self.path+"train_data.csv",sep=',',names=["id","image"])
        self.img1_transformer = img1_transformer
        self.img2_transformer = img2_transformer
        self.image_path = glob.glob(self.path+"images/*.jpg")
        self.len = len(self.image_path)
        self.image_name = self.data["image"]
        self.image_id = self.data["id"]
    def __len__(self):
        return self.len
    def __getitem__(self, index):
        index = index+1
        get_same_class = random.randint(0,1)
        img_1 = self.image_name[index]
        img_1_id = self.image_id[index]
        image1 = cv2.imread(self.path+'images/'+img_1, cv2.COLOR_BGR2RGB)
        augmented1 = self.img1_transformer(image = image1)
        label =get_same_class
        label = torch.tensor(label).unsqueeze(0)
        if get_same_class == 1:

            img_index = self.data[self.data["image"]==img_1].index.to_list()
            total_index_list = self.data[self.data["id"]==img_1_id].index.to_list()
            total_index_list.remove(img_index[0])
            img_2 = self.data["image"][random.choice(total_index_list)]
            image2 =cv2.imread(self.path+'images/'+img_2,cv2.COLOR_BGR2RGB)
            augmented2 = self.img2_transformer(image=image2)
            return augmented1["image"].float(),augmented2["image"].float(),label.float()
        if get_same_class == 0:

            total_index_list = self.data[self.data["id"] != img_1_id].index.to_list()
            #total_index_list.remove(0)
            #print(total_index_list)
            img_2 = self.data["image"][random.choice(total_index_list)]
            image2 = cv2.imread(self.path + 'images/' + img_2, cv2.COLOR_BGR2RGB)
            augmented2 = self.img1_transformer(image=image2)
            return augmented1["image"].float(),augmented2["image"].float(),label.float()

if __name__ == '__main__':






    #x = torch.randn(3,3)
    #y = torch.randn(3, 3)
    #data = pd.read_csv(Config.path + "train_data.csv", sep=',', names=["id", "image"])[0:2000]
    #print(len(data["id"]))
    #print(data['image'][train_idx])
    kfold = KFold(n_splits=5)
    dataset = Simdataset(Config.path,img1_transformer=img1_transformer,img2_transformer=img1_transformer)
    #dataloader = DataLoader(dataset,batch_size=2)
    # for i,(img1,img2,label) in enumerate(dataloader):
    #     print(img2.shape)

    #datloader = DataLoader()
    for i,(train_idx,val_idx) in enumerate(kfold.split(dataset)):
        #data = pd.read_csv(Config.path + "train_data.csv", sep=',', names=["id", "image"])[1:]
        #print(data["image"][train_idx])
        print(train_idx,val_idx)
       # print(val_idx.shape)
    #     dataset = Simdataset(Config.path,idx=train_idx)
    #     dataloader = DataLoader(dataset, batch_size=2)
    #
    #     # dataset = dataset[train_idx]
    #     # print(dataset)
    #     # dataloader = DataLoader(dataset,batch_size=2)
    #     for i,(img1,img2,label) in enumerate(dataloader):
    #         print(img2.shape)



